#include "analyzer_common.h"

// int analyzer_diagsolver_setup(
//     sparse_matrix_t* s, 
//     sparse_value_t* rhs, 
//     sparse_pointer_t* diag_ptr, 
//     sparse_index_t** nondiag_rowidx, 
//     sparse_index_t* nondiag_rowcnt,
//     sparse_index_t* nondiag_array_capacity
// ){
//     for(sparse_index_t row = 0; row < s->m; row++){
//         sparse_index_t row_nnz = 0;
//         for(sparse_pointer_t idx = s->rowpointer[row]; idx < s->rowpointer[row+1]; idx++){
//             sparse_index_t col = s->columnindex[idx];
//             if(row==col){
//                 diag_ptr[row] = idx;
//             }
//             if(s->value[idx] != 0){
//                 row_nnz++;
//             }
//         }
//         if(row_nnz > 1){
//             if((*nondiag_rowcnt) >= (*nondiag_array_capacity)){
//                 (*nondiag_rowidx) = realloc(*nondiag_rowidx, sizeof(sparse_index_t) *(*nondiag_array_capacity + 100));
//                 (*nondiag_array_capacity)+=100;
//             }
//             (*nondiag_rowidx)[*nondiag_rowcnt] = row;
//             (*nondiag_rowcnt)++;
//         }
//     }
// }

// int analyzer_diagsolver_solve(
//     sparse_matrix_t* s, 
//     sparse_value_t* rhs, 
//     sparse_pointer_t* diag_ptr, 
//     sparse_index_t* nondiag_rowidx, 
//     sparse_index_t nondiag_rowcnt
// ){
//     if(nondiag_rowcnt != 1){
//         printf("Not supported\n");
//         return;
//     }

//     // #pragma omp parallel for
//     for(sparse_index_t row = 0; row < nondiag_rowidx[0]; row++){
//         rhs[row] /= s->value[diag_ptr[row]];
//     }
//     // #pragma omp parallel for
//     for(sparse_index_t row = nondiag_rowidx[0]+1; row < s->m; row++){
//         rhs[row] /= s->value[diag_ptr[row]];
//     }

//     sparse_index_t diag_rowidx = nondiag_rowidx[0];
//     sparse_pointer_t diag_idx = diag_ptr[diag_rowidx];
//     for(sparse_pointer_t idx = s->rowpointer[diag_rowidx]; idx < diag_idx; idx++){
//         rhs[diag_rowidx] -= s->value[idx] * rhs[s->columnindex[idx]];
//     }
//     for(sparse_pointer_t idx = diag_idx+1; idx < s->rowpointer[diag_rowidx+1]; idx++){
//         rhs[diag_rowidx] -= s->value[idx] * rhs[s->columnindex[idx]];
//     }

//     rhs[diag_rowidx] /= s->value[diag_idx];
// }

int analyzer_diagsolver_setup(
    sparse_matrix_t* s, 
    sparse_value_t* rhs, 
    sparse_value_t* diag_value,
    sparse_index_t** nondiag_rowidx, 
    sparse_index_t* nondiag_rowcnt,
    sparse_index_t* nondiag_array_capacity
){
    for(sparse_index_t row = 0; row < s->m; row++){
        sparse_index_t row_nnz = 0;
        for(sparse_pointer_t idx = s->rowpointer[row]; idx < s->rowpointer[row+1]; idx++){
            sparse_index_t col = s->columnindex[idx];
            if(row==col){
                // diag_ptr[row] = idx;
                diag_value[row] = s->value[idx];
            }
            if(s->value[idx] != 0){
                row_nnz++;
            }
        }
        if(row_nnz > 1){
            if((*nondiag_rowcnt) >= (*nondiag_array_capacity)){
                (*nondiag_rowidx) = realloc(*nondiag_rowidx, sizeof(sparse_index_t) *(*nondiag_array_capacity + 100));
                (*nondiag_array_capacity)+=100;
            }
            (*nondiag_rowidx)[*nondiag_rowcnt] = row;
            (*nondiag_rowcnt)++;
        }
    }
}

int analyzer_diagsolver_solve(
    sparse_matrix_t* s, 
    sparse_value_t* rhs, 
    sparse_value_t* diag_value, 
    sparse_index_t* nondiag_rowidx, 
    sparse_index_t nondiag_rowcnt
){
    if(nondiag_rowcnt != 1){
        printf("Not supported\n");
        return;
    }

    // #pragma omp parallel for
    for(sparse_index_t row = 0; row < nondiag_rowidx[0]; row++){
        rhs[row] /= diag_value[row];
    }
    // #pragma omp parallel for
    for(sparse_index_t row = nondiag_rowidx[0]+1; row < s->m; row++){
        rhs[row] /= diag_value[row];
    }

    sparse_index_t diag_rowidx = nondiag_rowidx[0];
    // sparse_pointer_t diag_idx = diag_ptr[diag_rowidx];
    for(sparse_pointer_t idx = s->rowpointer[diag_rowidx]; idx < s->rowpointer[diag_rowidx+1]; idx++){
        if(s->columnindex[idx] != diag_rowidx){
            rhs[diag_rowidx] -= s->value[idx] * rhs[s->columnindex[idx]];
        }
    }

    rhs[diag_rowidx] /= diag_value[diag_rowidx];
}

// int analyzer_diagsolver_check(sparse_matrix_t* s, sparse_value_t* sol, sparse_value_t* rhs){
//     sparse_value_t* rhs_computed = malloc(sizeof(sparse_value_t) * s->m);
//     for(sparse_index_t row = 0; row < s->m; row++){
//         rhs_computed[row] = 0; 
//         for(sparse_pointer_t idx = s->rowpointer[row]; idx < s->rowpointer[row+1]; idx++){
//             rhs_computed[row] += s->value[idx] * sol[s->columnindex[idx]];
//         }
//     }
//     for(sparse_index_t row = 0; row < s->m; row++){
//         rhs_computed[row] -= rhs[row];
//     }
    
//     sparse_value_t norm2_computed = 0;
//     for(sparse_index_t row = 0; row < s->m; row++){
//         norm2_computed += (rhs_computed[row] * rhs_computed[row]);
//     }
//     norm2_computed = sqrt(norm2_computed);
    
//     sparse_value_t norm2_rhs = 0;
//     for(sparse_index_t row = 0; row < s->m; row++){
//         norm2_rhs += (rhs[row] * rhs[row]);
//     }
//     norm2_rhs = sqrt(norm2_rhs);

//     printf("[LYD solver] || Ax - b || / || b || = %le\n", norm2_computed / norm2_rhs);
// }

sparse_value_t analyzer_kahan_reduction(sparse_value_t* nums, sparse_pointer_t len) {
    sparse_value_t sum = 0.0;
    sparse_value_t c = 0.0;
    for (sparse_pointer_t i = 0; i < len; i++) {
        sparse_value_t y = nums[i] - c;
        sparse_value_t t = sum + y;
        c = (t - sum) - y;
        sum = t;
    }
    return sum;
}

int analyzer_diagsolver_check(sparse_matrix_t* s, sparse_value_t* sol, sparse_value_t* rhs){
    sparse_value_t* rhs_computed = malloc(sizeof(sparse_value_t) * s->m);
    sparse_value_t* tmp_kahan = malloc(sizeof(sparse_value_t) * s->m);
    memset(tmp_kahan, 0, sizeof(sparse_value_t) * s->m);
    sparse_index_t tmp_kahan_top = 0;
    for(sparse_index_t row = 0; row < s->m; row++){
        rhs_computed[row] = 0;
        memset(tmp_kahan, 0, sizeof(sparse_value_t) * tmp_kahan_top);
        tmp_kahan_top = 0;
        for(sparse_pointer_t idx = s->rowpointer[row]; idx < s->rowpointer[row+1]; idx++){
            // rhs_computed[row] += s->value[idx] * sol[s->columnindex[idx]];
            tmp_kahan[tmp_kahan_top] = s->value[idx] * sol[s->columnindex[idx]];
            tmp_kahan_top++;
        }
        rhs_computed[row] = analyzer_kahan_reduction(tmp_kahan, tmp_kahan_top);
    }
    for(sparse_index_t row = 0; row < s->m; row++){
        rhs_computed[row] -= rhs[row];
    }
    
    sparse_value_t norm2_computed = 0;
    for(sparse_index_t row = 0; row < s->m; row++){
        // norm2_computed += (rhs_computed[row] * rhs_computed[row]);
        rhs_computed[row] = (rhs_computed[row] * rhs_computed[row]);
    }
    norm2_computed = analyzer_kahan_reduction(rhs_computed, s->m);
    norm2_computed = sqrt(norm2_computed);
    
    sparse_value_t norm2_rhs = 0;
    for(sparse_index_t row = 0; row < s->m; row++){
        // norm2_rhs += (rhs[row] * rhs[row]);
        rhs[row] = (rhs[row] * rhs[row]);
    }
    norm2_rhs = analyzer_kahan_reduction(rhs, s->m);
    norm2_rhs = sqrt(norm2_rhs);

    printf("[LYD solver] (Kahan reduction) || Ax - b || / || b || = %le\n", norm2_computed / norm2_rhs);
}

int analyzer_diagsolver(sparse_matrix_t* s, sparse_value_t* rhs){

    sparse_value_t* diag_value = malloc(sizeof(sparse_value_t) * s->m);
    sparse_index_t* nondiag_rowidx = NULL;
    sparse_index_t nondiag_rowcnt = 0;
    sparse_index_t nondiag_array_capacity = 0;
    sparse_value_t* rhs_bkp = malloc(sizeof(sparse_value_t) * s->m);
    memcpy(rhs_bkp, rhs, sizeof(sparse_value_t) * s->m);

    struct timeval start, end;

    gettimeofday(&start, NULL);
    analyzer_diagsolver_setup(s, rhs, diag_value, &nondiag_rowidx, &nondiag_rowcnt, &nondiag_array_capacity);
    gettimeofday(&end, NULL);
    printf("[LYD solver] setup time : %.2lf ms\n", ((double)end.tv_sec - start.tv_sec) * 1e3 + ((double)end.tv_usec - start.tv_usec) * 1e-3);

    printf("[LYD solver] nondiag_rowcnt = %d, nondiag_rowidx[0] = %d\n", nondiag_rowcnt, nondiag_rowidx[0]);

    gettimeofday(&start, NULL);
    analyzer_diagsolver_solve(s, rhs, diag_value, nondiag_rowidx, nondiag_rowcnt);
    gettimeofday(&end, NULL);
    printf("[LYD solver] solve time : %.2lf ms\n", ((double)end.tv_sec - start.tv_sec) * 1e3 + ((double)end.tv_usec - start.tv_usec) * 1e-3);

    sparse_value_t x_min = 1e100;
    sparse_value_t x_max = 1e-100;
    for(sparse_index_t row = 0; row < s->m; row++){
        if(rhs[row] > x_max){
            x_max = rhs[row];
        }
        if(rhs[row] < x_min){
            x_min = rhs[row];
        }
    }

    printf("[LYD solver] x_min = %lf, x_max = %lf\n", x_min, x_max);

    analyzer_diagsolver_check(s, rhs, rhs_bkp);

    free(diag_value);
    free(nondiag_rowidx);
    
}